import matplotlib.pyplot as plt
import os, torch
import numpy as np

PYPLOT_PATH = 'output/plot.tar'


def get_init_data():
    if PYPLOT_PATH is not None and os.path.isfile(PYPLOT_PATH):
        tmp = torch.load(PYPLOT_PATH)
    tmp['train']['acc'] = sample(tmp['train']['acc'])
    tmp['train']['loss'] = sample(tmp['train']['loss'])
    return tmp

    # log_string("-> loaded checkpoint %s (epoch: %d)" % (CHECKPOINT_PATH, start_epoch))


def sample(list):
    # size=100
    list = np.array(list)
    list = list.reshape(-1,50)
    list = list.mean(1)
    return list


def create_graph(data):
    fig, axs = plt.subplots(1, 3)
    i = 0
    for key, value in data.items():
        my_plotter(axs[i], range(len(value)), value)
        axs[i].set_title(key)
        i = i + 1
    plt.show()


def my_plotter(ax, data1, data2, param_dict={}):
    """
    A helper function to make a graph

    Parameters
    ----------
    ax : Axes
        The axes to draw to

    data1 : array
       The x data

    data2 : array
       The y data

    param_dict : dict
       Dictionary of kwargs to pass to ax.plot

    Returns
    -------
    out : list
        list of artists added
    """
    out = ax.plot(data1, data2, **param_dict)
    return out


if __name__ == "__main__":

    plot_data = get_init_data()
    for _, data in plot_data.items():
        create_graph(data)
